import torch
from transformers import AutoModel, AutoTokenizer, AutoModelForSequenceClassification

model_path = 'D:/ai/huggingface-models/'

checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
# checkpoint = "bert-base-uncased"

raw_input = "I've been waiting for a HuggingFace course my whole life."

tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder=checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(model_path, subfolder=checkpoint)

tokens = tokenizer.tokenize(raw_input)
ids = tokenizer.convert_tokens_to_ids(tokens)

# 如果ids上不多加一个维度，model(input_ids)会报错
# input_ids = torch.tensor(ids)
# output = model(input_ids)

input_ids = torch.tensor([ids])
print(f'input ids:{input_ids}')

output = model(input_ids)
print(f'logits = {output.logits}')
